//you should define SIZEX/SIZEY/SIZEZ/BLKX/BLKY/ITER
//actually 2d-partitioning in this program 

#include <stdlib.h>
#include <stdio.h>
#include <math.h>
#include <assert.h>

#include <sys/time.h>
#include "3d7p_aceMesh.h"
#include "aceMesh_runtime_c.h"
#include "task.h"
double A[2][SIZEX][SIZEY][SIZEZ]; //attribute ((aligned(64)));
//for test
double test[2][SIZEX][SIZEY][SIZEZ];

extern void  slave_stencil_core(targs* args) ;
//arg
#define alpha_d 0.0876
#define beta_d  0.0765

void my_neighbors(void** neighbor_addrs,int* num_neigh, targs* args)
{
   int starti=args->starti;
    int startj=args->startj;
    void* src_addr=args->src_addr;
  
    *num_neigh=0;
    if(starti - BLKX>=0)
        neighbor_addrs[(*num_neigh)++]=(void*)(src_addr - BLKX * SIZEY * SIZEZ * 8);
    if(starti + BLKX < SIZEX)
        neighbor_addrs[(*num_neigh)++]=(void*)(src_addr + BLKX * SIZEY * SIZEZ * 8);
    if(startj - BLKY >= 0)
        neighbor_addrs[(*num_neigh)++]=(void*)(src_addr - BLKY * SIZEZ * 8);
    if(startj + BLKY < SIZEY)
        neighbor_addrs[(*num_neigh)++]=(void*)(src_addr + BLKY * SIZEZ* 8);
}


int Do3d7p(const int iter, const int thread_num)
{
    targs each;
    int itr,i,j;
     acemesh_begin_split_task("3d7p iter");
    for(i = 1; i < SIZEX -1; i += BLKX)
        {   
            int endi = (i + BLKX > SIZEX -1 ) ? SIZEX - 1 : i + BLKX;
            for(j = 1; j < SIZEY -1; j += BLKY)
            {   
                int endj = (j + BLKY > SIZEY -1 ) ? SIZEY - 1 : j + BLKY;
                each.starti=i;
                each.endi=endi;
                each.startj=j;
                each.endj=endj;
                each.k_N=SIZEZ;
                each.iter=itr;
                each.src_addr=&A[1][i][j][0];
                acemesh_push_wlist(1, &A[0][i][j][0],NORMAL);
                acemesh_push_rlist(1, &A[1][i][j][0], NORMAL);
                acemesh_task_generator_with_neighbors((TASK_FUNCPTR)slave_stencil_core,
                      (void*)(&each), sizeof(targs), 
                      NULL, (NEIGHBOR_FUNCPTR)my_neighbors, &each);
                acemesh_task_set_type(STENCIL_TASK);
            }   
      }
        acemesh_end_split_task();
    for(itr = 1; itr < iter; ++itr)
    {
        acemesh_begin_split_task("3d7p iter");
        for(i = 1; i < SIZEX -1; i += BLKX)
        {
            int endi = (i + BLKX > SIZEX -1 ) ? SIZEX - 1 : i + BLKX;
            for(j = 1; j < SIZEY -1; j += BLKY)
            {
                int endj = (j + BLKY > SIZEY -1 ) ? SIZEY - 1 : j + BLKY;
                each.starti=i;
                each.endi=endi;
                each.startj=j;
                each.endj=endj;
                each.k_N=SIZEZ;
                each.iter=itr;
                each.src_addr=&A[(itr+1)%2][i][j][0];
                acemesh_push_wlist(1, &A[itr%2][i][j][0],NORMAL);
                acemesh_push_rlist(1, &A[(itr+1)%2][i][j][0], NORMAL);
                acemesh_task_generator_with_neighbors((TASK_FUNCPTR)slave_stencil_core,
                      (void*)(&each), sizeof(targs), 
                      NULL, (NEIGHBOR_FUNCPTR)my_neighbors, &each);
                acemesh_task_set_type(STENCIL_TASK);
                acemesh_task_set_affinity(FOLLOW_AFFINITY);
            }
        }
        acemesh_end_split_task();
    }
    acemesh_spawn_and_wait(1);
    return 0;
}



void init()
{
    int i,j,k;
    for( i = 0; i < SIZEX; i++)
     for( j = 0; j < SIZEY; j++)
      for( k = 0; k < SIZEZ; k++){
        A[1][i][j][k] = (double)(i*2.5 + j*3.3 + k*0.5 ) /3.0; 
     }
}
int check(const int iter)
{
   int i,j,k;
   int itr;
   int ii,jj,kk;
   printf("begin check\n");
    for(i = 0; i < SIZEX; i++)
     for( j = 0; j < SIZEY; j++)
      for( k = 0; k < SIZEZ; k++){
         test[1][i][j][k] = (double)(i*2.5 + j*3.3 + k*0.5 ) /3.0; 
      }
       
    for(itr = 0; itr < iter; ++itr){
     for(i = 1; i < SIZEX - 1; i++)
      for( j = 1; j < SIZEY - 1; j++)
       for( k = 1; k < SIZEZ - 1; k++)
           test[itr%2][i][j][k] = alpha_d * test[(itr+1)%2][i][j][k] + \
                        beta_d * (test[(itr+1)%2][i-1][j][k] + test[(itr+1)%2][i+1][j][k] +\
                        test[(itr+1)%2][i][j-1][k] + test[(itr+1)%2][i][j+1][k] +\
                        test[(itr+1)%2][i][j][k-1] + test[(itr+1)%2][i][j][k+1] ); 
      }
       
    for(ii = 1; ii < SIZEX - 1; ii++)
      for(jj = 1; jj < SIZEY - 1; jj++)
       for(kk = 1; kk < SIZEZ - 1; kk++){
          if( fabs(test[(iter+1)%2][ii][jj][kk]  - A[(iter+1)%2][ii][jj][kk]) > 1e-13 ){
             printf("i: %d   j:  %d k: %d\n",ii,jj,kk);    
             printf("test: %lf",test[(iter+1)%2][ii][jj][kk]);
             printf("but A: %lf",A[(iter+1)%2][ii][jj][kk]);
             printf("they are different!\n");               
             return -1;
         }
      }
  printf("correct\n");
  return 0;
}



int main(int argc, char** argv)
{
    int ITER,num_threads;
    int i,j,k;
    struct timeval start;
    struct timeval end;
    double total_time;       
    int p = 1;
    ITER = 100;
    num_threads = 1;
    if(argc > 1)
        num_threads = atoi(argv[1]);
    if(argc > 2)
        ITER = atoi(argv[2]);
    init();
    gettimeofday(&start,NULL);
    //while(p);
    acemesh_runtime_init(num_threads);

    Do3d7p(ITER, num_threads);
    gettimeofday(&end,NULL);   
    acemesh_runtime_shutdown();  
#ifdef CHECK
    total_time = 0.0;
    total_time+=(end.tv_sec-start.tv_sec)+(end.tv_usec-start.tv_usec)/1000000.0;
    if(!check(ITER))
        printf("time: %lf\n",total_time);       
    else 
    {
        printf("check_error\n");
        exit(1);
    }
#endif
    return 0;
}
